import os
import json
import requests
import time
from openai import OpenAI
from utils.file_utils import save_to_file

class ImageGenerator:
    """
    图像生成器类，用于根据提示词生成图片
    """
    def __init__(self, config=None):
        """
        初始化图像生成器
        
        参数:
        config: 配置字典，包含API密钥等配置信息
        """
        # 如果没有提供配置，使用默认配置
        if config is None:
            config = self._load_default_config()
            
        # 保存配置
        self.config = config
        
        # 初始化OpenAI客户端
        self.client = self._create_client()
        
    def _load_default_config(self):
        """\加载默认配置"""
        # 尝试从配置文件加载配置
        config_path = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), 'config.json')
        try:
            with open(config_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        except Exception as e:
            print(f"加载配置文件失败: {e}")
            # 返回默认配置
            return {
                "api_keys": {
                    "openai": "your_openai_api_key_here"
                },
                "image_settings": {
                    "style": "cartoon",
                    "resolution": "1024x1024",
                    "number_of_images": 1
                }
            }
    
    def _create_client(self):
        """\创建OpenAI客户端"""
        # 获取API密钥
        api_key = self.config.get("api_keys", {}).get("openai", "")
        
        # 如果API密钥为空，尝试从环境变量获取
        if not api_key:
            api_key = os.environ.get("ARK_API_KEY")
            
        # 初始化OpenAI客户端
        return OpenAI(
            # 此为默认路径，您可根据业务所在地域进行配置
            base_url="https://ark.cn-beijing.volces.com/api/v3",
            # 从环境变量中获取您的 API Key
            api_key=api_key,
        )
    
    def generate_image_from_prompt_file(self, file_path="output/image_prompts.txt", prompt_index=0, download=False, save_dir="output"):
        """
        从文件中读取提示词并生成图片
        
        参数:
        file_path: 提示词文件路径
        prompt_index: 要使用的提示词索引，默认为0（第一个提示词）
        download: 是否下载图片，默认为False
        save_dir: 图片保存目录，默认为"output"
        
        返回:
        生成的图片URL或保存路径
        """
        try:
            # 读取提示词文件
            with open(file_path, 'r', encoding='utf-8') as f:
                prompts = f.readlines()
                
            # 检查索引是否有效
            if not prompts:
                raise ValueError(f"提示词文件 {file_path} 为空")
            
            if prompt_index < 0 or prompt_index >= len(prompts):
                raise ValueError(f"提示词索引 {prompt_index} 超出范围，有效范围为 0-{len(prompts)-1}")
            
            # 获取指定索引的提示词并去除行号和换行符
            raw_prompt = prompts[prompt_index].strip()
            # 移除行号前缀（如"1. "）
            prompt = raw_prompt.split('. ', 1)[1] if '. ' in raw_prompt else raw_prompt
            
            # 使用提示词索引作为文件名的一部分
            filename = f"image_prompt_{prompt_index + 1}"
            return self.generate_image(prompt, download=download, save_dir=save_dir, filename=filename)
            
        except FileNotFoundError:
            print(f"错误: 找不到提示词文件 {file_path}")
            raise
        except Exception as e:
            print(f"生成图片时发生错误: {str(e)}")
            raise
    
    def download_image(self, image_url, save_path):
        """
        下载图片并保存到指定路径
        
        参数:
        image_url: 图片的URL
        save_path: 保存图片的路径
        
        返回:
        保存成功的路径
        """
        try:
            # 发送HTTP请求获取图片数据
            response = requests.get(image_url, stream=True)
            response.raise_for_status()  # 检查请求是否成功
            
            # 确保保存目录存在
            save_dir = os.path.dirname(save_path)
            if save_dir and not os.path.exists(save_dir):
                os.makedirs(save_dir)
            
            # 保存图片数据
            with open(save_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    if chunk:
                        f.write(chunk)
            
            print(f"图片已保存到: {save_path}")
            return save_path
            
        except Exception as e:
            print(f"下载图片时发生错误: {str(e)}")
            raise
    
    def generate_image(self, prompt, download=False, save_dir="output", filename=None):
        """
        根据给定的提示词生成图片
        
        参数:
        prompt: 图片生成的提示词
        download: 是否下载图片，默认为False
        save_dir: 图片保存目录，默认为"output"
        filename: 保存的文件名，不包含扩展名，若为None则自动生成
        
        返回:
        生成的图片URL或保存路径
        """
        try:
            print(f"使用提示词生成图片: {prompt}")
            
            # 调用API生成图片
            images_response = self.client.images.generate(
                model="ep-20251001220231-h25st",
                prompt=prompt,
                size="2K",
                response_format="url",
                extra_body={
                    "watermark": True,
                },
            )
            
            # 获取生成的图片URL
            image_url = images_response.data[0].url
            print(f"图片生成成功，URL: {image_url}")
            
            # 如果需要下载图片
            if download:
                # 确保save_dir是绝对路径
                if not os.path.isabs(save_dir):
                    project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
                    save_dir = os.path.join(project_root, save_dir)
                
                # 生成文件名
                if filename is None:
                    # 使用时间戳和简短提示词作为文件名
                    timestamp = int(time.time())
                    short_prompt = prompt[:20].replace(' ', '_').replace('/', '_')
                    filename = f"image_{timestamp}_{short_prompt}"
                
                # 添加文件扩展名
                save_path = os.path.join(save_dir, f"{filename}.png")
                
                # 下载并保存图片
                return self.download_image(image_url, save_path)
            
            return image_url
            
        except Exception as e:
            print(f"生成图片时发生错误: {str(e)}")
            raise
    
    def generate_all_images_from_file(self, file_path="output/image_prompts.txt", download=False, save_dir="output"):
        """
        从文件中读取所有提示词并生成图片
        
        参数:
        file_path: 提示词文件路径
        download: 是否下载图片，默认为False
        save_dir: 图片保存目录，默认为"output"
        
        返回:
        生成的图片URL或保存路径列表
        """
        image_results = []
        try:
            # 读取提示词文件
            with open(file_path, 'r', encoding='utf-8') as f:
                prompts = f.readlines()
                
            # 为每个提示词生成图片
            for i, raw_prompt in enumerate(prompts):
                # 移除行号前缀和换行符
                prompt = raw_prompt.strip().split('. ', 1)[1] if '. ' in raw_prompt.strip() else raw_prompt.strip()
                
                # 生成图片并添加到列表
                filename = f"image_prompt_{i + 1}"
                result = self.generate_image(prompt, download=download, save_dir=save_dir, filename=filename)
                image_results.append((i + 1, result))  # 保存行号和结果（URL或路径）
            
            # 保存结果到文件
            results = [f"{i}. {res}" for i, res in image_results]
            save_to_file("\n".join(results), "image_urls.txt", "output")
            
            return image_results
            
        except Exception as e:
            print(f"生成所有图片时发生错误: {str(e)}")
            raise

# 用于测试的主函数
if __name__ == "__main__":
    try:
        # 初始化图像生成器
        image_generator = ImageGenerator()
        
        # 获取脚本所在目录，确保可以正确找到output目录
        script_dir = os.path.dirname(os.path.abspath(__file__))
        project_root = os.path.dirname(script_dir)  # 项目根目录
        prompt_file_path = os.path.join(project_root, "output", "image_prompts.txt")
        
        # 生成第一张图片
        print("生成第一张图片:")
        image_url = image_generator.generate_image_from_prompt_file(prompt_file_path)
        print(f"\n生成的图片URL: {image_url}")
        
    except Exception as e:
        print(f"程序执行失败: {str(e)}")